#include <iostream>
#include <vector>
#include <algorithm>
#include <stack>
using namespace std;

/*
隐马尔可夫模型：预测算法：维特比算法
*/

void hmm(double A[3][3], double B[3][2], double pai[], int O[], int k, int i) {
    //初始化prob[k][i]：在第k次观测为状态i的概率
    double prob[k][i];
    fill(prob[0], prob[0] + k*i, 0);
    // priv表示在第k次观测选取的是状态i，那么上一次观测可能性最高的选取状态为盒子priv[k][i]
    int priv[k][i];
    fill(priv[0], priv[0] + k*i, 0);
    for (int m = 0; m < i; m++) {
        //计算prob[0][m], m为遍历状态，第1次观测的结果为O[1],而第m个状态取得该结果的概率为B[m][O[0]]
        //第一次会选择第m个状态的概率为pai[m]
        prob[0][m] = pai[m] * B[m][O[0]];
    }
    
    //第2次....
    for (int c = 1; c < k; c++) {//c表示第c次,外循环遍历次数
        //计算第 c次，由状态为盒子 1,2,i 观测到 O[c]的概率
        for (int m = 0; m < i; m++) {//内循环遍历状态，盒子
            //本次计算的是 prob[c][m],计算的是第m个盒子
            int argmax = 0;  //argmax记录的是盒子的编号
            double psi = 0;

            for (int n = 0; n < i; n++) {//
                if (prob[c-1][n] * A[n][m] * B[m][O[c]] > prob[c][m])
                    prob[c][m] = prob[c-1][n] * A[n][m] * B[m][O[c]];
                //cout << prob[c-1][n] * A[n][m] << endl;
                if (prob[c-1][n] * A[n][m] > psi) {
                    argmax = n;
                    psi = prob[c-1][n] * A[n][m];
                }
            }

            priv[c][m] = argmax;
        }
    }

    //那么最优路径的终点为：max(prob[k-1])对应的盒子状态
    double temp_max = 0;
    int argmax = 0;
    for (int m = 0; m < i; m++) {
        //找到最优路径终点对应的状态
        if (prob[k-1][m] > temp_max) {
            temp_max = prob[k-1][m];
            argmax = m;
        }
    }
    stack<int> s;
    s.push(argmax);
    //由最终路径的终点逆向找到其他点
    for (int c = k-1; c >= 1; c--) {
        s.push(priv[k-1][argmax]);
        argmax = priv[k-1][argmax];
    }
    //输出结果
    while (!s.empty()) {
        cout << s.top() << endl;
        s.pop();
    }

}
int main() {

    // A:状态转移概率矩阵
    double A[3][3] = {0.5, 0.2, 0.3, 0.3, 0.5, 0.2, 0.2, 0.3, 0.5};
    // B:观测概率矩阵
    double B[3][2] = {0.5, 0.5, 0.4, 0.6, 0.7, 0.3};
    // pai:初始选择状态的概率矩阵
    double pai[3] = {0.2, 0.4, 0.4};
    // O观测序列, 0表示观测类别0，1表示观测类别1
    int O[3] = {0, 1, 0};
    // k观测次数, i可选状态的个数
    int k = 3, i = 3;
    hmm(A, B, pai, O, k, i);
    return 0;
}